import torch
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.transforms.functional as trf
import os
import PIL
from PIL import Image


_default_dataset_dir = os.path.join(os.path.dirname(__file__), "../dataset/anime-face/cropped")


class Dataset(torch.utils.data.Dataset):
    def __init__(self, output_height=32, output_width=32, dir_path:str=_default_dataset_dir):
        super().__init__()
        self.output_height = output_height
        self.output_width = output_width
        files = os.listdir(dir_path)
        self.transform = transforms.Compose([
            transforms.Resize((self.output_height, self.output_width)),
            transforms.ToTensor()
        ])
        self.file_list = []
        for f in files:
            f = os.path.join(dir_path, f)
            if not f.endswith('.jpg'):
                print("Find a file which is not a jpg, ", f)
                continue
            self.file_list.append(f)
        print("{} files registered.".format(len(self.file_list)))

    """
    override!
    return a tensor of shape (3, 128, 128)
    """
    def __getitem__(self, item):
        f = self.file_list[item]
        try:
            img = Image.open(f)
            if img.width < self.output_width or img.height < self.output_height:  # 略过太小的图片
                return self.__getitem__((item + 1) % self.__len__())
        except PIL.UnidentifiedImageError as e:  # 略过无法读取的图片
            # print("file({}) damaged, skip it.".format(f))
            return self.__getitem__((item + 1) % self.__len__())
        tensor = self.__preprocess(img)
        return tensor

    # override
    def __len__(self):
        return len(self.file_list)

    def __preprocess(self, image_pil):
        return self.transform(image_pil)


def main():
    d = Dataset()
    loader = torch.utils.data.DataLoader(d, 2, True)
    for data in loader:
        print(data)
        print(data.shape)
        t1 = data[0, :]
        t2 = data[1, :]
        import matplotlib.pyplot as plt
        plt.subplot(121)
        plt.imshow(t1.permute(1, 2, 0).numpy())
        plt.subplot(122)
        plt.imshow(t2.permute(1, 2, 0).numpy())
        plt.show()


if __name__ == '__main__':
    main()
